package com.mega.common.core.utils;


import com.mega.common.core.web.model.TreeNode;
import lombok.experimental.UtilityClass;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * Tree工具类
 */
@UtilityClass
public class TreeUtil {

    /**
     * 创建树结构
     *
     * @param treeNodes
     * @return
     */
    public <T extends TreeNode> List<T> buildTree(List<T> treeNodes, String rootCode) {
        List<T> rootNodes = new ArrayList<>();
        Iterator<T> it = treeNodes.iterator();
        while (it.hasNext()) {
            T next = it.next();
            if (rootCode.equals(next.getParentCode())) {
                rootNodes.add(next);
                it.remove();
            }
        }
        for (T regionTree : rootNodes) {
            List<T> child = getChild(treeNodes, regionTree.getCode().toString());
            regionTree.setChildren(child);
        }
        return rootNodes;
    }

    /**
     * 查子节点
     *
     * @param treeNodes
     * @param parentCode
     * @return
     */
    private <T extends TreeNode> List<T> getChild(List<T> treeNodes, String parentCode) {
        //子节点列表
        List<T> childList = new ArrayList<>();
        Iterator<T> it = treeNodes.iterator();
        while (it.hasNext()) {
            T treeNode = it.next();
            if (parentCode.equals(treeNode.getParentCode())) {
                childList.add(treeNode);
                it.remove();
            }
        }

        //遍历 递归获取子节点的子节点
        for (T treeNode : childList) {
            List<T> child = getChild(treeNodes, treeNode.getCode().toString());
            treeNode.setChildren(child);
        }
        //递归出口  childList长度为0
        if (childList.size() == 0) {
            return new ArrayList<>();
        }
        return childList;
    }
}
